from typing import (
    List, 
    Union, 
    Optional,
    Any,
    Dict
)

import json
import os

from langchain_core.utils.function_calling import (
    convert_to_openai_function,
)
from loguru import logger

from langchain_core.tools import tool
from langchain_deepseek import ChatDeepSeek
from langchain_core.tools.structured import StructuredTool

class FunctionCallingWrapper:
    def __init__(self, functions: List[StructuredTool]):
        """
        初始化函数调用包装器
        
        Args:
            llm: 大语言模型实例
            functions: 函数字典，key为函数名，value为可调用对象
        """
        self.llm = self._create_llm()
        self.functions = dict([(f.name, f.func) for f in functions])
        self.function_descs = [f"{convert_to_openai_function(tool, strict=True)}" for tool in functions]

    def _create_llm(self) -> ChatDeepSeek:
        openai_api_base = os.environ["OPENAI_CHAT_ENDPOINT"]
        model = os.environ["OPENAI_CHAT_MODEL"]
        key = os.environ["OPENAI_CHAT_KEY"]

        chatopenai = ChatDeepSeek(
            temperature=0.0,
            api_base=openai_api_base,
            model=model,
            api_key=key,
            streaming=False,
            max_tokens=3096,
        )
        return chatopenai

    def _build_prompt_with_functions_history(
        self, 
        input: str, 
        history: List[dict]
    ) -> str:
        """
        构建包含函数信息的提示词
        
        Args:
            input: 原始问题
            
        Returns:
            包含函数信息的提示词
        """
        function_descriptions = "\n".join(self.function_descs)
        
        if len(history) > 0:
            # 包含其它工具调用结果
            scratchpad = f"已经执行过的API及其结果如下JSON描述：\n{history}\n 请理解上述语义后按照接下来的步骤思考"
        else:
            scratchpad = ""

        return f"""
-- context --
{scratchpad}
-- end context --

1. 如果根据上述上下文信息，可以直接输出最终答案，请**立即**按照如下JSON格式输出最终结果：
{{
    "result": "The final answer ..."
}}

2. 如果上述信息不足，请仔细分析下面用JSON描述的可用API工具名字、功能及其参数和返回值：
-- API List --
{function_descriptions} 
-- end API List --

请结合输入问题推理是否调用一个或者多个以上列表中的API。如果需要，请严格遵循以下规则：
1）**尽量减少API调用的次数**，仅在必要时调用，避免重复（在context中）或不必要的调用。
2）函数名称和参数名称必须严格匹配API List中的定义，**不能自行臆造任何不在API List中的函数调用**。
3）请确保所有调用的函数和参数名称完全符合API List中的定义，包括大小写和命名方式，**不要使用任何自定义的名称**。
4）请严格按照如下示例输出JSON格式内容：
{{
    "api_call": [{{
        "function": "function name",
        "parameters": {{
            "param1": "value1",
            "param2": "value2"
            ....
        }}
    }}]
}}

3. 用户输入请求为：
{input}

4. 结合context和步骤1/2/3中的信息，思考回答用户问题的解决步骤，并输出接下来*一个步骤*应该执行的操作：
"""
    
    def _parse_function_call(self, text: str) -> Optional[Dict[str, Any]]:
        """
        解析大模型返回的函数调用信息
        
        Args:
            response: 大模型返回的响应
            
        Returns:
            解析后的函数调用信息，如果不需要函数调用则返回None
        """
        try:
            from langchain_core.utils.json import parse_json_markdown
            # 获得</think>标签之后的文本数据
            pos = text.find("</think>")
            if pos > 0:
                text = text[pos + 8:]

            text = text.strip()
            # Markdown 格式输出的JSON格式解析
            return parse_json_markdown(text)
        except json.JSONDecodeError:
            pass
        return None
    
    def _execute_function(self, function_call: Dict[str, Any]) -> Union[str, Dict[str, Any]]:
        """
        执行函数调用
        
        Args:
            function_call: 函数调用信息
            
        Returns:
            函数执行结果
        """
        for f in function_call["api_call"]:
            func_name = f["function"]
            if func := self.functions.get(func_name):
                params = f.get("parameters", {})
                f["result"] = f"{func(**params)}"
            else:
                raise ValueError(f"Function {func_name} not found")    
            
        return function_call
    
    def run(self, input: str) -> str:
        """
        运行函数调用流程
        
        Args:
            input: 原始提示词
            
        Returns:
            最终结果
        """
        historical_calls = []
        
        while True:
            current_prompt = self._build_prompt_with_functions_history(input, historical_calls)

            # 调用大模型
            response = self.llm(current_prompt)
            
            # 解析函数调用
            partial = self._parse_function_call(response.content)
            
            if "result" in partial:
                # 不需要函数调用，直接返回结果
                return partial.get("result", "")
                
            # 执行函数
            partial = self._execute_function(partial)
            
            # 更新函数调用历史
            historical_calls.extend(partial["api_call"])

def test_function_calling():
    @tool
    def get_current_datetime() -> str:
        """获得当前的系统日期和时间信息"""
        return "当前的日期和时间为：2025/02/14 13:03:09" #假设现在是春天

    @tool        
    def get_current_weather(city: str) -> int:
        """获得指定城市的当前的天气情况"""
        if city=="哈尔滨":
            return "哈尔滨当前的气温-23摄氏度，大雪，风速1.4米/秒"
        elif city=="海口":
            return "海口当前的气温15摄氏度，晴朗，风速2.1米/秒"
        elif city=="上海":
            return "上海当前的气温8摄氏度，晴朗，风速0.8米/秒"        
        else:
            return f"{city}当前的气温45摄氏度，晴朗，风速8.8米/秒"

    wrapper = FunctionCallingWrapper([get_current_weather, get_current_datetime])
    
    # 测试函数调用
    result = wrapper.run("有一个用户春夏季节喜欢在北方（比如哈尔滨、沈阳、青岛）度假，秋冬喜欢在南方（比如上海、杭州或者海口、福州等）度假，请根据当前的时间建议该用户去哪里度假，并根据天气推荐穿什么类型的衣服")
    assert isinstance(result, str)
    assert "哈尔滨" in result
